import argparse
import train_util as util
from bert_config import bert_net_cfg
import mindspore as ms
import bert_model as bert
import train_util_2 as util2
import eval_util as eval


def parse_args():
    # 创建解析
    parser = argparse.ArgumentParser(description="classifier config",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # 添加参数
    parser.add_argument('--train_data_path', type=str, default='data/train.csv')
    parser.add_argument('--test_data_path', type=str, default='data/test.csv')
    parser.add_argument('--vocab_path', type=str, default='data/vocab.txt')
    parser.add_argument('--max_seq_len',type=int,default=128)
    parser.add_argument('--batch_size',type=int,default=8)
    parser.add_argument('--train_ratio',type=float,default=0.8)
    parser.add_argument('--out_path',type=str,default='inference/model')
    parser.add_argument('--loss',type=str,default='softmaxcrossentropywithlogis')
    parser.add_argument('--fc_dim1',type=int,default=300)
    parser.add_argument('--class_num',type=int,default=5)
    parser.add_argument('--learning_rate',type=float,default=0.001)
    parser.add_argument('--epochs',type=int,default=10)
    parser.add_argument('--bert_weight',type=str,default='inference/model/bert.ckpt')
    parser.add_argument('--model_weight',type=str,default='inference/model/finetune-4_2500.ckpt')
    parser.add_argument('--glove_path',type=str,default='data/glove.6B.300d.txt')
    parser.add_argument('--device',type=str,default='gpu')
    parser.add_argument('--bert_vocab_size',type=int,default=21128)
    parser.add_argument('--type_vocab_size',type=int,default=2)
    parser.add_argument('--hidden_size',type=int,default=256)
    parser.add_argument('--num_layers',type=int,default=2)
    parser.add_argument('--dropout',type=float,default=0.5)
    parser.add_argument('--model_type',type=str,default='lstm')
    parser.add_argument('--glove_dimention',type=int,default=300)
    parser.add_argument('--train_or_eval',type=str,default='train_and_eval')
    parser.add_argument('--result_path',type=str,default='data/result_example.txt')
    parser.add_argument('--split_ration',type=float,default=0.7)
    parser.add_argument('--split_or_not',type=int,default=0)

    # 解析参数
    args_opt = parser.parse_args()
    return args_opt

if __name__ == '__main__':
    args_opt = parse_args()
    if args_opt.device=='gpu':
        ms.set_context(mode=ms.PYNATIVE_MODE, device_target="GPU")
    elif args_opt.device=='npu':
        ms.set_context(mode=ms.PYNATIVE_MODE,device_target="Ascend")
    else:
        ms.set_context(mode=ms.PYNATIVE_MODE)
    if args_opt.train_or_eval=='train_and_eval' or args_opt.train_or_eval=='train':
        if args_opt.model_type=='bert':
            net=util.train(args_opt,bert.BertConfig(vocab_size=args_opt.bert_vocab_size,type_vocab_size=args_opt.type_vocab_size))
        else:
            net=util2.train(args_opt)
        if args_opt.train_or_eval=='train_and_eval':
            if args_opt.model_type=='bert':
                eval.test(args_opt,bert.BertConfig(vocab_size=args_opt.bert_vocab_size,type_vocab_size=args_opt.type_vocab_size),net)
            else:
                eval.test_tiny_model(args_opt,net)
    elif args_opt.train_or_eval=='eval':
        eval.test(args_opt,bert.BertConfig(vocab_size=args_opt.bert_vocab_size,type_vocab_size=args_opt.type_vocab_size))